package cn.frank.framework.caesar.module.ai.openAi.conf.arrangement;

import cn.frank.framework.caesar.module.ai.openAi.conf.arrangement.assistant.Assistant;
import cn.frank.framework.caesar.module.ai.openAi.conf.arrangement.assistant.AssistantCall;
import cn.frank.framework.caesar.module.ai.openAi.conf.arrangement.assistant.AssistantUnique;
import cn.frank.framework.caesar.module.ai.openAi.conf.arrangement.function.ToolsService;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

@Slf4j
@Component()
@Order(value = -1)
public class OpenAiConfig {

    /**
     * 创建一个内存中的 EmbeddingStore，用于存储文本向量。
     * 可用于知识库检索（RAG）功能。
     */
    @Bean
    public EmbeddingStore embeddingStore() {
        return new InMemoryEmbeddingStore();
    }

    /**
     * 构建一个具有聊天能力的助手对象
     *
     * @param chatModel
     * @param streamingChatModel
     * @return
     */
    @Bean
    public Assistant assistant(ChatLanguageModel chatModel, StreamingChatLanguageModel streamingChatModel) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .streamingChatLanguageModel(streamingChatModel)
                .chatMemory(chatMemory)
                .build();
        return assistant;
    }

    /**
     * 构建一个具有带记忆功能的助手对象
     *
     * @param chatModel
     * @param streamingChatModel
     * @return
     */
    @Bean
    public AssistantUnique assistantUnique(ChatLanguageModel chatModel, StreamingChatLanguageModel streamingChatModel) {

        AssistantUnique assistant = AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(chatModel)
                .streamingChatLanguageModel(streamingChatModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder().maxMessages(10).id(memoryId).build())
                .build();
        return assistant;
    }

    /**
     * 构建一个具有回调功能的助手对象
     * @param chatModel
     * @param streamingChatModel
     * @param toolsService
     * @return
     */
    @Bean
    public AssistantCall assistantCall(ChatLanguageModel chatModel, StreamingChatLanguageModel streamingChatModel, ToolsService toolsService) {

        AssistantCall assistant = AiServices.builder(AssistantCall.class)
                .tools(toolsService)
                .chatLanguageModel(chatModel)
                .streamingChatLanguageModel(streamingChatModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder().maxMessages(10).id(memoryId).build())
                .build();
        return assistant;
    }

}
